// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <melon/container/f14_set.h>
#include <pollux/common/memory/raw_vector.h>
#include <pollux/vector/base_vector.h>

namespace kumo::pollux {
    struct VectorValueSetEntry {
        const BaseVector *vector;
        vector_size_t index;
    };

    struct VectorValueSetHasher {
        size_t operator()(const VectorValueSetEntry &entry) const {
            return entry.vector->hash_value_at(entry.index);
        }
    };

    struct VectorValueSetComparer {
        bool operator()(
            const VectorValueSetEntry &left,
            const VectorValueSetEntry &right) const {
            return left.vector->equal_value_at(right.vector, left.index, right.index);
        }
    };

    using VectorValueSet = melon::F14FastSet<
        VectorValueSetEntry,
        VectorValueSetHasher,
        VectorValueSetComparer>;

    /// A map translating values in a vector to positions in an alphabet
    /// vector. if the values are complex type or variable length, also
    /// keeps track of the serialized size of the values. Usable for
    /// ad-hoc reencoding for serialization.
    class VectorMap {
    public:
        // Constructs 'this' to index the distinct elements in 'alphabet'. 'alphabet'
        // is not changed and not owned. For example, when concatenating two
        // dictionaries, the first initializes the map.
        explicit VectorMap(BaseVector &alphabet);

        // Constructs an empty map initializing alphabet to an empty vector of 'type'.
        // Alphabet is owned. 'reserve' is the expected count of distinct values.
        VectorMap(
            const TypePtr &type,
            memory::MemoryPool *pool,
            int32_t reserve = 32);

        /// Assigns a zero-based id to each distinct value in 'vector' at positions
        /// 'rows'. The ids are returned in 'ids'. If insertToAlphabet is true and the
        /// value is not previously in alphabet_, the value is added to it. Alphabet
        /// must be owned if insertToAlphabet is true.
        void addMultiple(
            BaseVector &vector,
            melon::Range<const vector_size_t *> rows,
            bool insertToAlphabet,
            vector_size_t *ids);

        // Gets/assigns s an id to single value. The meaning of parameters is as in
        // addMultiple().
        vector_size_t addOne(
            const BaseVector &vector,
            vector_size_t row,
            bool insertToAlphabet = true);

        /// Returns the number of distinct values in 'this'.
        vector_size_t size() const {
            return isString_
                       ? distinctStrings_.size() + (nullIndex_ != kNoNullIndex)
                       : distinctSet_.size();
        }

        /// Returns the approximate serialized binary length of the 'index'th value in
        /// 'alphabet_'. The type must be a variable length type.
        vector_size_t lengthAt(vector_size_t index) const {
            POLLUX_DCHECK_EQ(fixedWidth_, kVariableWidth);
            return alphabetSizes_[index];
        }

        const VectorPtr &alphabetOwned() const {
            return alphabetOwned_;
        }

    private:
        static constexpr vector_size_t kNoNullIndex = -1;
        static constexpr int32_t kVariableWidth = -1;

        // Vector containing all the distinct values.
        BaseVector *alphabet_;

        // Optional owning pointer to 'alphabet_'.
        VectorPtr alphabetOwned_;

        // Map from value in 'alphabet_' to the index in 'alphabet_'.
        VectorValueSet distinctSet_;

        // Map from string value in 'alphabet_' to index in 'alphabet_'. Used only if
        // 'alphabet_' is a FlatVector<StringView>.
        melon::F14FastMap<StringView, int32_t> distinctStrings_;

        // Index of null value in 'alphabet_'.
        vector_size_t nullIndex_{kNoNullIndex};

        // Serialized size estimate for the corresponding element of 'alphabet'.
        raw_vector<vector_size_t> alphabetSizes_;

        // True if  using 'distinctStrings_'
        const bool isString_;

        const int32_t fixedWidth_;
    };
} // namespace kumo::pollux
